use super::tool::*;
use anyhow::{anyhow, Result};
use serde::{Deserialize, Serialize};
use strum::{Display, EnumMessage, EnumVariantNames};

#[derive(Debug, Clone, Serialize, Display, EnumVariantNames, EnumMessage)]
#[serde(rename_all = "snake_case", tag = "role")]
pub enum ChatMessage {
    /// A message from a system.
    System(SystemMessage),
    /// A message from a human.
    User(UserMessage),
    /// A message from the assistant.
    Assistant(AssistantMessage),
    /// A message from a tool.
    Tool(ToolMessage),
}

#[derive(Debug, Clone, Serialize)]
pub struct SystemMessage {
    /// The contents of the system message.
    content: String,
    /// An optional name for the participant. Provides the model information to differentiate between participants of the same role.
    #[serde(skip_serializing_if = "Option::is_none")]
    name: Option<String>,
}

#[derive(Debug, Clone, Serialize)]
pub struct UserMessage {
    /// The contents of the user message.
    content: String,
    /// An optional name for the participant. Provides the model information to differentiate between participants of the same role.
    #[serde(skip_serializing_if = "Option::is_none")]
    name: Option<String>,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AssistantMessage {
    /// The contents of the system message.
    #[serde(default)]
    pub content: Option<String>,
    /// An optional name for the participant. Provides the model information to differentiate between participants of the same role.
    #[serde(skip_serializing_if = "Option::is_none", default)]
    pub name: Option<String>,    
    /// The tool calls generated by the model, such as function calls
    #[serde(skip_serializing_if = "Vec::is_empty", default)]
    pub tool_calls: Vec<ToolCall>,
}

#[derive(Debug, Clone, Serialize)]
pub struct ToolMessage {
    /// The contents of the tool message.
    content: String,
    /// Tool call that this message is responding to.
    tool_call_id: String,
}

impl ChatMessage {
    pub fn system(content: impl Into<String>) -> Self {
        Self::System(SystemMessage {
            content: content.into(),
            name: None,
        })
    }

    pub fn user(content: impl Into<String>) -> Self {
        Self::User(UserMessage {
            content: content.into(),
            name: None,
        })
    }

    pub fn assistant(content: Option<impl Into<String>>, tool_calls: impl Into<Vec<ToolCall>>) -> Self {
        Self::Assistant(AssistantMessage {
            content: content.map(|c| c.into()),
            name: None,
            tool_calls: tool_calls.into(),
        })
    }

    pub fn tool(content: impl Into<String>, tool_call_id: impl Into<String>) -> Self {
        Self::Tool(ToolMessage {
            content: content.into(),
            tool_call_id: tool_call_id.into()
        })
    }

    pub fn as_system(&self) -> Result<&SystemMessage> {
        if let Self::System(msg) = self {
            Ok(msg)
        } else {
            Err(anyhow!("Except system message!"))
        }
    }

    pub fn as_user(&self) -> Result<&UserMessage> {
        if let Self::User(msg) = self {
            Ok(msg)
        } else {
            Err(anyhow!("Except user message!"))
        }
    }

    pub fn as_assistant(&self) -> Result<&AssistantMessage> {
        if let Self::Assistant(msg) = self {
            Ok(msg)
        } else {
            Err(anyhow!("Except assistant message!"))
        }
    }

    pub fn as_tool(&self) -> Result<&ToolMessage> {
        if let Self::Tool(msg) = self {
            Ok(msg)
        } else {
            Err(anyhow!("Except system message!"))
        }
    }
}
